import torch.nn as nn
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops

class GNNMessagePassing(MessagePassing):
    def __init__(self, hidden=16):
        super().__init__(aggr='add')
        self.mlp = nn.Sequential(
            nn.Linear(2, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1)   # 输出关联强度
        )

    def forward(self, x, edge_index):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j 邻居特征 → 过 MLP → 消息
        return self.mlp(x_j)

    def update(self, aggr_out):
        return aggr_out
